import os
import random
import shutil

folder_path = r"/media/champrin/Data/Desktop/RoboMaster-Unity/Screenshots"
save_path = r"/home/champrin/Desktop/MV-CS016-10UC+DA1041860/给预备役/Screenshots split"
os.makedirs(save_path, exist_ok=True)

folder_train = os.path.join(save_path, "train")  # os.path.dirname得到上一级目录
os.makedirs(folder_train, exist_ok=True)
folder_val = os.path.join(save_path, "val")  # os.path.dirname得到上一级目录
os.makedirs(folder_val, exist_ok=True)
folder_test = os.path.join(save_path, "test")  # os.path.dirname得到上一级目录
os.makedirs(folder_test, exist_ok=True)

# 指定比例 a:b:c
a = 7.5
b = 2.5
c = 0

# 获取文件夹A中的所有图片文件
image_files = [f for f in os.listdir(folder_path) if f.endswith('.png')]

# 随机打乱文件列表
random.shuffle(image_files)

# 计算每个分组的数量
total_files = len(image_files)
a_count = int(total_files * (a / (a + b + c)))
b_count = int(total_files * (b / (a + b + c)))
c_count = total_files - a_count - b_count

# 将文件分配到不同文件夹中
for idx, file_name in enumerate(image_files):
    if idx < a_count:
        shutil.copy(os.path.join(folder_path, file_name), folder_train)
    elif a_count <= idx < a_count + b_count:
        shutil.copy(os.path.join(folder_path, file_name), folder_val)
    else:
        shutil.copy(os.path.join(folder_path, file_name), folder_test)

print(f"Files distributed to train: {a_count}, val: {b_count}, test: {c_count}")